#include <algorithm>
#include "geometryflow.h"
#include "../contour/mcconextractor.h"

using namespace PBVP;

GeometryFlow::GeometryFlow(Reg3DVolumeData* phi, Reg3DVolumeData *func)
{
	mp_Phi = phi;
	mp_Func = func;

	unsigned int dim[3];
	float span[3];
	mp_Func->getDimension(dim);
	mp_Func->getSpan(span);
	nx = dim[0]; ny = dim[1]; nz = dim[2];
	dx = span[0]; dy = span[1]; dz = span[2];

	mp_Up = new float[nx*ny*nz];

	m_MaxDepth = 4;
}

GeometryFlow::~GeometryFlow()
{
	delete [] mp_Up;
}

void GeometryFlow::evolve(int maxIterations)
{
	mp_Func->calcGradient();
	double CFLMax = 0;
	
	for (unsigned int k = 0; k < nz; k++)
	{
		for (unsigned int j = 0; j < ny; j++)
		{
			for (unsigned int i = 0; i < nx; i++)
			{
				double grad[3];
				mp_Func->getVertGrad(i, j, k, grad);
				grad[0] = (fabs(grad[0]) > 20)? 20:grad[0];
				grad[1] = (fabs(grad[1]) > 20)? 20:grad[1];
				grad[2] = (fabs(grad[2]) > 20)? 20:grad[2];
				double cond = fabs(grad[0])/dx + fabs(grad[1])/dy + fabs(grad[2])/dz;
				if(CFLMax < cond) CFLMax = cond;

			}
		}
	}

	// Time increment for CFL stability
	// m_dt = std::min<double>(0.9 / CFLMax, 0.015);
	m_dt = 0.9 / CFLMax;
	printf("Time increment = %f\n", m_dt);

	int nit = 0;
	while (nit < maxIterations)
	{
		for (int i = 0; i < nx*ny*nz; i++)
		{
			// truncate the function phi
			if (mp_Phi->getValue(i) > m_MaxDepth)
			{
				mp_Phi->setValue(m_MaxDepth, i);
			} else if (mp_Phi->getValue(i) < -m_MaxDepth)
			{
				mp_Phi->setValue(-m_MaxDepth, i);
			}
		}

		flow();

		if (nit > 0 && (nit%10 == 0))
		{
			reinitilization();
		}

		nit ++;
		// testing purpose
		// printf("iteration # %d\n", nit);
	}

}

void GeometryFlow::flow()
{
	for (unsigned int k = 0; k < nz; k++)
	{
		for (unsigned int j = 0; j < ny; j++)
		{
			for (unsigned int i = 0; i < nx; i++)
			{
				mp_Up[k*nx*ny+j*nx+i] = 0;
				if (fabs(mp_Phi->getValue(i, j, k)) < m_MaxDepth)
				{
					// compute the 1st and 2nd derivatives of the level set function Phi
					double phi_grad[3], phi_par2[6];
					mp_Phi->getVertGrad(i, j, k, phi_grad);
					mp_Phi->getVertParDeriv2(i, j, k, phi_par2);
					double grad2 = phi_grad[0]*phi_grad[0] + phi_grad[1]*phi_grad[1] + phi_grad[2]*phi_grad[2];
					if(grad2 < 0.0001) grad2 = 0.0001;		// Is 0.0001 a good truncation
					double grad1 = sqrt(grad2);
					// Mean curvature
					double H;
					H = -phi_grad[0]*(phi_par2[0]*phi_grad[0]+phi_par2[1]*phi_grad[1]+phi_par2[2]*phi_grad[2])
						-phi_grad[1]*(phi_par2[1]*phi_grad[0]+phi_par2[3]*phi_grad[1]+phi_par2[4]*phi_grad[2])
						-phi_grad[2]*(phi_par2[2]*phi_grad[0]+phi_par2[4]*phi_grad[1]+phi_par2[5]*phi_grad[2]);
					H = H/grad2 + phi_par2[0] + phi_par2[3] + phi_par2[5];		// 

					// Get the gradient of driving function at (i, j, k)
					double func_grad[3];
					mp_Func->getVertGrad(i, j, k, func_grad);

					// upwind 1st derivative of the function Phi
					// Godunov's scheme
					double forward[3], back[3];		// forward and backward derivatives
					//mp_Phi->getVertGradForward(i, j, k, forward);
					//mp_Phi->getVertGradBackward(i, j, k, back);
					mp_Phi->getQuadENODerivatives(i, j, k, forward, back);
					double upwind[3];

					for (int l = 0; l < 3; l++)
					{
						if(func_grad[l]*forward[l] > 0 && func_grad[l]*back[l] >0) {
							upwind[l] = back[l];
						} else if (func_grad[l]*forward[l] < 0 && func_grad[l]*back[l] < 0) {
							upwind[l] = forward[l];
						} else if (func_grad[l]*back[l] <= 0 && func_grad[l]*forward[l] >= 0) {
							upwind[l] = 0;
						} else {	// func_grad[l]*back[l] >= 0 && func_grad[l]*forward[l] <= 0
							upwind[l] = (fabs(back[l]) >= fabs(forward[l]))? back[l]:forward[l];
						}
					}

					mp_Up[k*nx*ny+j*nx+i] = func_grad[0]*upwind[0] + func_grad[1]*upwind[1] + func_grad[2]*upwind[2] +
						std::min<double>(mp_Func->getValue(i, j, k), 0.01) * H;
					//mp_Up[k*nx*ny+j*nx+i] = -sqrt(upwind[0]*upwind[0] + upwind[1]*upwind[1] + upwind[2]*upwind[2]);
				}
			}
		}
	}

	for (int i = 0; i < nx*ny*nz; i++)
	{
		if (fabs(mp_Phi->getValue(i)) < m_MaxDepth)
		{
			double val = mp_Phi->getValue(i) + mp_Up[i]*m_dt;
			mp_Phi->setValue(val, i);
		}
	}
}

void GeometryFlow::reinitilization()
{
	printf("reinitialization ...\n");
	for (int it = 0; it < 20; it++)
	{
		for (int k = 0; k < nz; k++)
		{
			for (int j = 0; j < ny; j++)
			{
				for (int i = 0; i < nx; i++)
				{
					mp_Up[k*nx*ny+j*nx+i] = 0;

					if (mp_Phi->getOneNeighborAbsMin(i, j, k) < m_MaxDepth)
					{
						double plus[3], minus[3];
						mp_Phi->getQuadENODerivatives(i, j, k, plus, minus);

						double grad2 = 0.5*(minus[0]*minus[0]+minus[1]*minus[1]+minus[2]*minus[2] +
							plus[0]*plus[0]+plus[1]*plus[1]+plus[2]*plus[2]);
						double grad1 = sqrt(grad2);
						// Osher's book p. 67 
						double phi = mp_Phi->getValue(i, j, k); 
						double sd = phi / sqrt(phi*phi + grad2*dx*dx);
						double sdp = (sd < 0.0)? 0 : sd;
						double sdm = (sd > 0.0)? 0 : sd;

						double offset = 0;
						if(sdp > 0) {	// sd is positive
							// Rouy and Tourin version of Godunov's method
							double dx2 = std::max<double>(std::max<double>(minus[0], 0.0)*std::max<double>(minus[0], 0.0), 
								std::min<double>(plus[0], 0.0)*std::min<double>(plus[0], 0.0));
							double dy2 = std::max<double>(std::max<double>(minus[1], 0.0)*std::max<double>(minus[1], 0.0),
								std::min<double>(plus[1], 0.0)*std::min<double>(plus[1], 0.0));
							double dz2 = std::max<double>(std::max<double>(minus[2], 0.0)*std::max<double>(minus[2], 0.0),
								std::min<double>(plus[2], 0.0)*std::min<double>(plus[2], 0.0));
							double len = sqrt(dx2+dy2+dz2);
							mp_Up[k*nx*ny+j*nx+i] = - sdp * (len - 1);
						} else if (sdm < 0)
						{
							double dx2 = std::max<double>(std::min<double>(minus[0], 0.0)*std::min<double>(minus[0], 0.0), 
								std::max<double>(plus[0], 0.0)*std::max<double>(plus[0], 0.0));
							double dy2 = std::max<double>(std::min<double>(minus[1], 0.0)*std::min<double>(minus[1], 0.0), 
								std::max<double>(plus[1], 0.0)*std::max<double>(plus[1], 0.0));
							double dz2 = std::max<double>(std::min<double>(minus[2], 0.0)*std::min<double>(minus[2], 0.0), 
								std::max<double>(plus[2], 0.0)*std::max<double>(plus[2], 0.0));
							double len = sqrt(dx2 + dy2 + dz2);
							mp_Up[k*nx*ny+j*nx+i] = -sdm * (len - 1);
						}
					}
				}
			}
		}
		double dt = std::min<double>(0.5*dx, 0.05*m_MaxDepth);
		for (int v = 0; v < nx*ny*nz; v++)
		{
			double val = mp_Phi->getValue(v) + dt*mp_Up[v];
			mp_Phi->setValue(val, v);
			if (mp_Phi->getValue(v) > m_MaxDepth)
			{
				mp_Phi->setValue(m_MaxDepth, v);
			} else if (mp_Phi->getValue(v) < -m_MaxDepth)
			{
				mp_Phi->setValue(-m_MaxDepth, v);
			}
		}
	}
}

